'''
The All Pairs Shortest Path (APSP) calculates 
the shortest (weighted) path between all pairs of nodes.
https://neo4j.com/docs/graph-data-science/current/algorithms/all-pairs-shortest-path/
'''

from priorityQueue import PriorityQueue
from graph import Graph
import sys,csv,unittest


def dijkstra(G,start):
    pq = PriorityQueue()
    start.setDistance(0)
    pq.buildHeap([(v.getDistance(),v) for v in G])
    while not pq.isEmpty():
        nowVert = pq.delMin()
        for nextVert in nowVert.getConnections():
            newDist = nowVert.getWeight(nextVert) + nowVert.getDistance()
            if newDist < nextVert.getDistance():
                nextVert.setDistance(newDist)
                nextVert.setPred(nowVert)
                pq.decreaseKey(nextVert,newDist)

def getAllPairsShortestPath(G,keys):
    result = [['#' for _ in range(len(keys))]for _ in range(len(keys))]
    for key in keys:
        dijkstra(G,G.getVertex(key))
        for v in G:
            if v.getDistance() not in (0,sys.maxsize):
                G.addEdge(key,v.getId(),v.getDistance())
            v.setDistance(sys.maxsize)
    
    for i,key in enumerate(keys):
        Vertfrom = G.getVertex(key)
        for neighbor in Vertfrom.getConnections():
            j = keys.find(neighbor.getId())
            result[i][j] = str(Vertfrom.getWeight(neighbor))
    
    return result
    

class TestAPSP(unittest.TestCase):
    def setUp(self):
        self.G = Graph()
    
    def test_case1(self):
        self.G.addEdge("A","B",1)
        self.G.addEdge("A","E",1)
        self.G.addEdge("B","C",1)
        self.G.addEdge("B","F",1)
        self.G.addEdge("C","G",1)
        self.G.addEdge("D","C",1)
        self.G.addEdge("E","F",1)
        self.G.addEdge("E","I",1)
        self.G.addEdge("F","J",1)
        self.G.addEdge("G","B",1)
        self.G.addEdge("G","K",1)
        self.G.addEdge("H","D",1)
        self.G.addEdge("H","K",1)
        self.G.addEdge("J","I",1)
        self.G.addEdge("J","G",1)
        self.G.addEdge("J","K",1)
        self.G.addEdge("K","L",1)
        self.G.addEdge("L","H",1)
        answer = ['#12712362345','##16#1253234','#2#5#3145423',
                  '#31##4256534','#456#1351234','#345##242123',
                  '#124#2#34312','#421#53#7612','############',
                  '#234#3131#12','#643#75298#1','#532#641872#']
        result = getAllPairsShortestPath(self.G,"ABCDEFGHIJKL")
        for line1,line2 in zip(answer,result):
            self.assertEqual(line1,''.join(line2))
    
    def test_case2(self):
        self.G.addEdge("A","B",50)
        self.G.addEdge("A","C",50)
        self.G.addEdge("A","D",100)
        self.G.addEdge("B","D",40)
        self.G.addEdge("C","D",40)
        self.G.addEdge("C","E",80)
        self.G.addEdge("D","E",30)
        self.G.addEdge("D","F",80)
        self.G.addEdge("E","F",40)
        answer = [['#', '50', '50', '90', '120', '160'],
                  ['#', '#', '#', '40', '70', '110'],
                  ['#', '#', '#', '40', '70', '110'],
                  ['#', '#', '#', '#', '30', '70'],
                  ['#', '#', '#', '#', '#', '40'],
                  ['#', '#', '#', '#', '#', '#']]
        result = getAllPairsShortestPath(self.G,"ABCDEF")
        for line1,line2 in zip(answer,result):
            self.assertEqual('\t'.join(line1),'\t'.join(line2))

if __name__ == "__main__":
    unittest.main()
